#!/usr/bin/env python3
"""
扫描并（可选）修复 annotation txt 文件：

每行格式示例：
0_0535.jpg:
0_0536.jpg: 3 0.009766 0.311806 0.019531 0.145833

规则：每一行在冒号后至多保留 5 个 token（类别 + 4 个坐标），如果出现更多（表示有多个物体），截断为只保留第一个物体。

脚本功能：
- 递归查找指定目录下的 .txt 文件
- 检测并统计存在问题的行（冒号后 token 数 > 5）
- 打印存在问题的文件路径和问题行数总计
- 询问用户是否修复；若确认，备份原文件并修改文件（截断多余 token）

用法示例：
python truncate_annotation_lines.py -d E:\\datasets\\Huawei

"""
import argparse
import os
import sys
import fnmatch
from collections import defaultdict

# 全局默认设置（可在脚本顶部直接修改）
# 如果希望硬编码扫描目录，在这里把 DEFAULT_ROOT 设置为类似 r'E:\\datasets\\Huawei'
# 否则保持 None，使用命令行参数或当前目录。
DEFAULT_ROOT =  '/home_ssd/lhc/hand_detect_v3' # r"E:\datasets\Huawei\output_2cam_2_test"  # e.g. r'E:\\datasets\\Huawei' or None

# 默认要匹配的文件名或 glob 模式，例如:
#  - 'annotations.txt' 精确文件名
#  - '*.txt' 匹配所有 txt 文件（默认）
#  - 'train*.txt' 支持通配符
DEFAULT_FILENAME = "event_frame_window_new.txt" # "rgb_wrap_window_new.txt"  # e.g. 'annotations.txt' or None to use '*.txt'


def analyze_file(path):
    """返回 (problem_lines, new_lines) 其中 problem_lines 是问题行的索引集合（0-based）和原行内容，
    new_lines 是修改后的所有行（仅在需要修复时使用）。"""
    problem_lines = []
    new_lines = []
    with open(path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    for idx, line in enumerate(lines):
        stripped = line.rstrip('\n')
        if ':' not in stripped:
            # 没有冒号的行不做处理
            new_lines.append(line)
            continue
        head, tail = stripped.split(':', 1)
        tail_stripped = tail.strip()
        if tail_stripped == '':
            # 没有标注
            new_lines.append(line)
            continue
        tokens = tail_stripped.split()
        if len(tokens) > 5:
            # 问题行：多于一个对象（每个对象占 5 个 token）
            problem_lines.append((idx, line))
            # 生成新的行内容，仅保留前 5 个 token
            kept = ' '.join(tokens[:5])
            new_line = f"{head}: {kept}\n"
            new_lines.append(new_line)
        else:
            new_lines.append(line)

    return problem_lines, new_lines


def scan_directory(root, filename_pattern='*.txt'):
    """递归扫描目录，返回问题文件映射和总问题行数

    filename_pattern 可以是精确文件名或 glob 模式（例如 '*.txt'）。
    """
    problem_map = {}  # path -> list of (idx, orig_line)
    total = 0
    for dirpath, dirnames, filenames in os.walk(root):
        for fn in filenames:
            # 使用 fnmatch 对文件名进行匹配，允许传入精确名或通配符
            if not fnmatch.fnmatch(fn, filename_pattern):
                continue
            path = os.path.join(dirpath, fn)
            problems, _ = analyze_file(path)
            if problems:
                problem_map[path] = problems
                total += len(problems)
    return problem_map, total


def backup_file(path):
    bak = path + '.bak'
    if not os.path.exists(bak):
        os.rename(path, bak)
        return bak
    else:
        # 如果 .bak 存在，创建带时间戳的备份
        import time
        ts = int(time.time())
        bak2 = f"{path}.bak.{ts}"
        os.rename(path, bak2)
        return bak2


def fix_files(problem_map):
    """对 problem_map 中的每个文件进行修复，先备份再写入新内容，返回已修复文件数和已修复的总行数"""
    fixed_files = 0
    fixed_lines = 0
    for path in problem_map:
        # 分析并获取 new_lines
        _, new_lines = analyze_file(path)
        # 备份
        bak = backup_file(path)
        # 写回
        with open(path, 'w', encoding='utf-8') as f:
            f.writelines(new_lines)
        fixed_files += 1
        fixed_lines += len(problem_map[path])
        print(f"已备份 {path} -> {bak}，并写入修复后的内容。")
    return fixed_files, fixed_lines


def main():
    parser = argparse.ArgumentParser(description='扫描并可选修复每行多于一个目标的 annotation txt 文件')
    parser.add_argument('-d', '--dir', default='.', help='要扫描的根目录（递归），默认当前目录，或使用脚本顶部 DEFAULT_ROOT')
    parser.add_argument('-n', '--name', default=None, help="要匹配的文件名或 glob 模式，例如 'annotations.txt' 或 '*.txt'（CLI 优先于脚本顶部 DEFAULT_FILENAME）")
    args = parser.parse_args()

    # 决定根目录：命令行参数优先，其次是脚本顶部的 DEFAULT_ROOT，再次默认为当前目录
    root = args.dir if args.dir != '.' or DEFAULT_ROOT is None else DEFAULT_ROOT
    if root is None:
        root = '.'
    if not os.path.isdir(root):
        print(f'错误：指定的路径不是目录：{root}', file=sys.stderr)
        sys.exit(2)

    # 决定文件名匹配模式：命令行参数优先，其次是脚本顶部的 DEFAULT_FILENAME，再次默认为 '*.txt'
    filename_pattern = args.name if getattr(args, 'name', None) else DEFAULT_FILENAME
    if not filename_pattern:
        filename_pattern = '*.txt'

    print(f"开始扫描目录：{root}（递归，文件模式：{filename_pattern}）...")
    problem_map, total = scan_directory(root, filename_pattern=filename_pattern)

    if not problem_map:
        print('未发现问题行。')
        return

    print('\n发现以下存在问题的文件（每行冒号后 token 数 > 5）：')
    for p, lst in problem_map.items():
        print(f'{p}  —— 问题行数: {len(lst)}')

    print(f'\n总共有 {len(problem_map)} 个文件包含问题，总问题行数：{total}')

    # 强制交互：始终询问用户是否修复（不支持跳过交互）
    ans = input('是否要对这些文件进行修复（会备份原文件为 <file>.bak）？ [y/N]: ').strip().lower()
    do_fix = ans == 'y' or ans == 'Y'

    if not do_fix:
        print('未执行修复，脚本退出。')
        return

    fixed_files, fixed_lines = fix_files(problem_map)
    print(f'修复完成：{fixed_files} 个文件，{fixed_lines} 条行被修复。')


if __name__ == '__main__':
    main()
